# https://www.cnblogs.com/for-technology-lover/p/14838525.html
import keras
import numpy as np

n = 3
# depth = n * 9 + 1
depth = n * 6 + 2


def resnet_layer(inputs, num_filters=16, kernel_size=3, strides=1, activation='relu', batch_normalization=True,
                 conv_first=True):
    # 使用何恺明初始化，l2正则
    conv = keras.layers.Conv2D(num_filters, kernel_size=kernel_size, padding='same', kernel_initializer='he_normal',
                               kernel_regularizer=keras.regularizers.l2(1e-4))
    x = inputs
    if conv_first:
        x = conv(x)
        if batch_normalization:
            x = keras.layers.BatchNormalization()(x)
        if activation is not None:
            x = keras.layers.Activation(activation)(x)
    else:
        if batch_normalization:
            x = keras.layers.BatchNormalization()(x)
        if activation is not None:
            x = keras.layers.Activation(activation)(x)

        x = conv(x)
    return x


def resnet(input_shape, depth, num_classes=10):
    if (depth - 2) % 6 != 0:
        raise ValueError('depth should be 6n+2')

    num_filters = 16
    num_res_blocks = int((depth - 2) / 6)

    inputs = keras.layers.Input(shape=input_shape)
    x = resnet_layer(inputs=inputs)
    for stack in range(3):
        for res_block in range(num_res_blocks):
            strides = 1
            if stack > 0 and res_block == 0:
                strides = 2

            y = resnet_layer(inputs=x, num_filters=num_filters, strides=strides)
            y = resnet_layer(inputs=y, num_filters=num_filters, activation=None)
            if stack > 0 and res_block == 0:
                x = resnet_layer(inputs=x, num_filters=num_filters, kernel_size=1, strides=strides, activation=None,
                                 batch_normalization=False)
            x = keras.layers.add([x, y])
            x = keras.layers.Activation('relu')(x)
        num_filters *= 2
    x = keras.layers.AveragePooling2D(pool_size=8)(x)
    x = keras.layers.Flatten()(x)
    outputs = keras.layers.Dense(num_classes, activation='softmax', kernel_initializer='he_normal')(x)
    model = keras.Model(inputs=inputs, outputs=outputs)
    return model


# input_shape = (224, 224, 3)
input_shape = (32, 32, 3)
model = resnet(input_shape=input_shape, depth=depth)
(x_train, y_train), (x_test, y_test) = keras.datasets.cifar10.load_data()
# 计算类别数
num_labels = len(np.unique(y_train))

# 转化为one-hot编码
y_train = keras.utils.to_categorical(y_train)
y_test = keras.utils.to_categorical(y_test)

input_shape = x_train.shape[1:]
x_train = x_train.astype('float32') / 255.
x_test = x_test.astype('float32') / 255.

batch_size = 64
# epochs = 200
epochs = 2
# 编译模型，使用分类交叉熵损失
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['acc'])
model.summary()
# 模型训练
model.fit(x_train, y_train, batch_size=batch_size, epochs=epochs, validation_data=(x_test, y_test), shuffle=True)
# 模型训练
scores = model.evaluate(x_test, y_test, batch_size=batch_size, verbose=0)
print('Test loss:', scores[0])
print('Test accuracy: ', scores[1])

print(444)
